package painting

import (
	"encoding/json"
	"gxcvu-ai/pkg/response"
	"io"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
)

type Wangxiang struct {
}

const (
	api_key = "sk-zA7xbusaoI"
)

type WangxiangInput struct {
	Prompt         string `json:"prompt"`
	NegativePrompt string `json:"negative_prompt,omitempty"`
}

type WangxiangParamenter struct {
	Style string `json:"style,omitempty"`
	Size  string `json:"size,omitempty"`
	Seed  int64  `json:"seed,omitempty"`
	Num   int64  `json:"n,omitempty"`
}

type WangxiangRequest struct {
	Model      string              `json:"model"`
	Input      WangxiangInput      `json:"input"`
	Parameters WangxiangParamenter `json:"parameters,omitempty"`
}

type WangxiangResponse struct {
	RequestId string `json:"request_id"`
	Output    struct {
		TaskId     string `json:"task_id"`
		TaskStatus string `json:"task_status"`
	} `json:"output"`
}

type WangxiangJobResponse struct {
	RequestId string `json:"request_id"`
	Output    struct {
		TaskId     string `json:"task_id"`
		TaskStatus string `json:"task_status"`
		Results    []struct {
			Url string `json:"url"`
		} `json:"results,omitempty"`
		TaskMetrics struct {
			Total   int `json:"TOTAL"`
			Success int `json:"SUCCEEDED"`
			Failed  int `json:"FAILED"`
		}
	} `json:"output"`
	Usage struct {
		ImageCount int `json:"image_count"`
	}
}

func (e Wangxiang) SubmitJob(prompt string, negative_prompt string, style string, size string, c *gin.Context) {
	const url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"

	payload := WangxiangRequest{
		Model: "wanx-v1",
		Input: WangxiangInput{
			Prompt:         prompt,
			NegativePrompt: negative_prompt,
		},
		Parameters: WangxiangParamenter{
			Style: style,
			Size:  size,
			Num:   1,
		},
	}

	postForm, _ := json.Marshal(payload)
	req, _ := http.NewRequest("POST", url, strings.NewReader(string(postForm)))
	req.Header.Add("Authorization", "Bearer "+api_key)
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("X-DashScope-Async", "enable")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		response.Error(c, err)
	}
	defer resp.Body.Close()
	data, _ := io.ReadAll(resp.Body)
	var jsonData WangxiangResponse
	err = json.Unmarshal(data, &jsonData)
	if err != nil {
		response.Error(c, err)
	}
	response.CreatedJSON(c, jsonData)
}

func (e Wangxiang) GetTaskStatus(taskId string, c *gin.Context) {
	url := "https://dashscope.aliyuncs.com/api/v1/tasks/" + taskId
	req, _ := http.NewRequest("GET", url, nil)
	req.Header.Add("Authorization", "Bearer "+api_key)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		response.Error(c, err)
	}
	defer resp.Body.Close()
	data, _ := io.ReadAll(resp.Body)
	var jsonData WangxiangJobResponse
	err = json.Unmarshal(data, &jsonData)
	if err != nil {
		response.Error(c, err)
	}
	response.JSON(c, jsonData)
}
